#############################################
#Robustness check: CATE for monotonic assignment orders
#############################################

rm(list = ls())

#Functions
source("functions.R")

#Packages
installPackageNotFound("stargazer")
installPackageNotFound("lmtest")
installPackageNotFound("sandwich")

#############################################
#Read in Voter data
#############################################

#All voters
pair_df <- read.csv("voter_pairs.csv", stringsAsFactors = F)

#Drop same household and alternate design volunteers
pair_df <- pair_df[which(pair_df$same_household == 0 &
                           pair_df$Volunteer_In_District == 1), ]

#############################################
#Calculate pair distance order
#############################################

#Sort by distance to pair
pair_df <- pair_df[order(pair_df$volunteer_id,
                         pair_df$average_distance_miles), ]

#Create index along volunteer-voter distance
pair_df$distance_index <- ave(pair_df$pair_id, pair_df$volunteer_id,
                              FUN = seq_along)

#Cap distance index at 4+
pair_df$distance_factor <- ifelse(pair_df$distance_index > 3, "4",
                                  pair_df$distance_index)

#Flag monotonic assignments
pair_df$is_monotonic <- ifelse(pair_df$assignment_factor ==
                                 pair_df$distance_factor, 1, 0)

#Flag first pair the closest
pair_df$first_closest <- ifelse(pair_df$assignment_factor == 1 &
                                  pair_df$distance_factor == 1, 1, 0)

#Index to count pairs
pair_df$pair_index <- 1

#Collapse by volunteer
volunteer_summary <- aggregate(cbind(is_monotonic, first_closest,
               pair_index)~volunteer_id, data = pair_df, sum)

#Get strictly monotonic volunteers
strict_monotonic <- volunteer_summary$volunteer_id[
  which(volunteer_summary$is_monotonic == volunteer_summary$pair_index)]

#Get first pair closest
first_closest <- volunteer_summary$volunteer_id[
  which(volunteer_summary$first_closest == 1)]

#############################################
#Paired treatment effect
#############################################
#Model variables
vars <- c("BinaryTurnout", "Black", "Hispanic", "Asian", "Male", "Under30",
          "VotePropensity", "Partisanship")

#Data groups
data_list <- list("all" = which(!is.na(pair_df$volunteer_id)),
  "monotonic" = which(pair_df$volunteer_id %in% strict_monotonic),
  "first_closest" = which(pair_df$volunteer_id %in% first_closest))

#Specifications
itt_base <- paste0("BinaryTurnout_diff ~1")
itt_controls <- paste0("BinaryTurnout_diff ~",
                       paste0(vars[2:length(vars)], "_diff", collapse = "+"))
itt_assignment <- paste0("BinaryTurnout_diff~ -1 + factor(assignment_factor)")
itt_assignment_controls <- paste0(itt_assignment, "+",
                                  paste0(vars[2:length(vars)], "_diff",
                                  collapse = "+"))
#List of formulas
formula_list <- c(itt_base,
                  itt_controls,
                  itt_assignment,
                  itt_assignment_controls)

#Empty list to hold model results
model_list <- list()
se_list <- list()

for(k in 1:length(data_list)){
  for(i in 1:length(formula_list)){
    #Subset to volunteer group
    this_group <- pair_df[data_list[[k]], ]

    #Calculate within-pair differences
    for(this_var in vars){
      this_group[, paste0(this_var, "_diff")] <-
        this_group[, paste0(this_var, "_treatment")] -
        this_group[, paste0(this_var, "_control")]
    }

    #Treatment effect at the pair level
    this_model <- lm(formula(formula_list[[i]]), data = this_group)

    #Robust SEs
    this_se <- coeftest(this_model, vcov = vcovHC(this_model, type = "HC3"))

    #Store results
    model_list[[names(data_list)[k]]][[i]] <- this_model
    se_list[[names(data_list)[k]]][[i]] <- this_se
  }
}

#############################################
#CATE Latex table
#############################################
title <- "CATE Models"

#Full sample model
latex_output <- stargazer(model_list[["all"]][[4]],
                          model_list[["first_closest"]][[4]],
                          model_list[["monotonic"]][[4]],
                          se = list(se_list[["all"]][[4]][, "Std. Error"],
                                    se_list[["first_closest"]][[4]][, "Std. Error"],
                                    se_list[["monotonic"]][[4]][, "Std. Error"]
                          ),
                          omit = paste0(vars[2:length(vars)], "_diff"),
                          omit.stat = c("ser", "rsq", "f"),
                          covariate.labels = c("Order 1", "Order 2", "Order 3",
                                               "Order 4+"),
                          add.lines = list(c("Voter Covariates",
                                             "\\multicolumn{1}{c}{Y}",
                                             "\\multicolumn{1}{c}{Y}",
                                             "\\multicolumn{1}{c}{Y}")),
                          column.separate = c(1,1,1,1),
                          dep.var.labels = "",
                          dep.var.caption = "",
                          notes.append = FALSE,
                          column.sep.width = "0pt",
                          no.space = TRUE,
                          title = title,
                          star.cutoffs = c(0.05, 0.01, 0.001),
                          type = "text")
